import torch
import numpy as np

from PIL import Image
from matplotlib import pyplot as plt

from models.facenet import FaceNet
from utils import letterbox_image

if __name__ == '__main__':
	input_shape = [160, 160, 3]
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

	model = FaceNet(num_classes=4)
	model.to(device)

	model_path = "xavier-logs/Epoch-100-Total_Loss-0.3660912911097209.pth-Val_Loss-0.3047718480229378.pth"
	model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
	model.eval()

	while True:
		image_1 = input("Input image_1 filename:")
		try:
			image_1 = Image.open(image_1)
		except Exception as e:
			print(e)
			print("Image_1 Open Error! Try again!")
			continue

		image_2 = input("Input image_2 filename:")
		try:
			image_2 = Image.open(image_2)
		except Exception as e:
			print(e)
			print("Image_2 Open Error! Try again!")
			continue

		with torch.no_grad():
			image_1 = letterbox_image(image_1, [input_shape[1], input_shape[0]])
			image_2 = letterbox_image(image_2, [input_shape[1], input_shape[0]])

			photo_1 = torch.from_numpy(
				np.expand_dims(np.transpose(np.asarray(image_1).astype(np.float64) / 255, (2, 0, 1)), 0)).type(
				torch.FloatTensor)
			photo_2 = torch.from_numpy(
				np.expand_dims(np.transpose(np.asarray(image_2).astype(np.float64) / 255, (2, 0, 1)), 0)).type(
				torch.FloatTensor)

			photo_1 = photo_1.to(device)
			photo_2 = photo_2.to(device)

			output1 = model(photo_1).cpu().numpy()
			output2 = model(photo_2).cpu().numpy()

			l1 = np.linalg.norm(output1 - output2, axis=1)

		plt.subplot(1, 2, 1)
		plt.imshow(np.array(image_1))

		plt.subplot(1, 2, 2)
		plt.imshow(np.array(image_2))
		plt.text(-12, -12, 'Distance:%.3f' % l1, ha='center', va='bottom', fontsize=11)
		plt.show()
